using DataFrames
using QuadGK
using Distributions
using CSV
using LinearAlgebra
using JLD
using Interpolations
using SparseArrays
using DelimitedFiles

#clearconsole()

#-------------------------------------------------------------
# IRF Configuration
#-------------------------------------------------------------

H = 40 # maximum horizon
n_every = 10 # use every n_every'th draw from the posterior
sh_size = 3   # shock size, in multiples of standard deviations
sh_id = 1 # sh_id = 1: TFP shock, sh_id = 2: GDP shock, sh_id = 3: Employment shock


#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
# include(readDir *"vech.jl");
include(readDir *"logSpline_Procedures.jl");
# include(readDir *"VAR_Procedures.jl");
include(readDir *"Loaddata.jl");
include(readDir *"VAR_MoreLags_Procedures.jl")

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nfVARSpec = "10tc"
nModSpec  = "1" # K=10
#nModSpec  = "2" # K=22
#nModSpec    = "3" # K=10, lambda2=lambda3=1.0
nMCMCSpec = "1"
modName   = "SS"  # VAR or SS

specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/" * modName * "spec" * nModSpec * ".jl")
include(specDir * "/" * modName * "MCMCspec" * nMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/"
agg_data, period_agg, mean_unrate = loadaggdata(SampleStart,SampleEnd)
n_agg    = size(agg_data)[2]


#-------------------------------------------------------------
# Load coefficients from density estimation and knots
#-------------------------------------------------------------
sNameLoadDir = "fVAR" * nfVARSpec
loaddir      = "$(pwd())/results/" * sNameLoadDir *"/";
sNameFile    = "K" * string(K) * "_fVAR" * nfVARSpec

PhatDensCoef_factor, MDD_GoF, VinvLam_all, period_Dens_ind, PhatDensCoef_lambda, PhatDensCoef_mean, PhatDensCoef_mean_allt  = loaddensdata(SampleStart,SampleEnd,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

knots_all = readdlm(loaddir * "fVAR" *nfVARSpec * "_knots_all.csv", ',', Float64, '\n'; skipstart = 1);
ii=getindex.(findall(K_vec.-K.==0),[1 2])[1] # find index ii where K==K_vec
knots = knots_all[quant_sel[ii,:].==1]

# knots=loadknots(quant_vec, quant_sel)

#-------------------------------------------------------------
# Adjust sample
#-------------------------------------------------------------

# Arrange data for VAR analysis
# We start by dropping the initial Tdrop-nlags observations
agg_data_adj = agg_data[Tdrop-nlags+1:end,:]
PhatDensCoef_factor_adj = PhatDensCoef_factor[Tdrop-nlags+1:end,:]
# We now have T*=(T-T0)+nlags observations and call the first observation t=1
# Observations t=1,...,nlags are used to initialize lags
# and observations t=nlags+1,...,T* are the estimation sample
# Note that by holding Tdrop fixed and changing nlags<=Tdrop we can change
# the nunber of lags holding the estimation sample fixed.
# Estimation sample has T*-nlags=T-T0 observations

# Define some constants
n       = n_agg+n_cross
p       = nlags

# Combine aggregate and cross sectional observations and create YY XX for VAR estimation
data_adj = [ agg_data_adj PhatDensCoef_factor_adj ]
YY, XX, PHIols, Sols, SIGMAols = data_to_YY(data_adj, nlags, nlags)
# FS: having nlags twice as input arguments is not necessary. We can simplify
# but would have to adjust other programs as well.



#-------------------------------------------------------------
# Load posterior draws
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_" * modName * nModSpec * "_" * "MCMC" * nMCMCSpec;
loadDir  = "$(pwd())/results/" * sName *"/";

Bpmean     = load(loadDir * sName * "_PostMeans.jld", "Bpmean")
Sigpmean   = load(loadDir * sName * "_PostMeans.jld", "Sigpmean")
Sigtrpmean = load(loadDir * sName * "_PostMeans.jld", "Sigtrpmean")

Bpdraw     = load(loadDir * sName * "_PostDraws.jld", "Bpdraw")
Sigpdraw   = load(loadDir * sName * "_PostDraws.jld", "Sigpdraw")
Sigtrpdraw = load(loadDir * sName * "_PostDraws.jld", "Sigtrpdraw")

#-------------------------------------------------------------
# Generate IRFs at posterior mean
#-------------------------------------------------------------

println("")
println("Generating IRFs at posterior mean... ")
println("")

YY_IRF          = zeros(H+1,n)
L               = Sigtrpmean
Btilde          = reshape(Bpmean, n*p,n)
response        = IRFredu(Btilde,L,H,sh_size,sh_id)
YY_IRF[2:end,:] = response'

PhatDens_IRF    = zeros(H+1,length(xgrid));
# compute density IRFs
for hh = 1:H+1
    PhatDensCoef_hh    = coefRecover(YY_IRF[hh,n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF[hh,3]+mean_unrate,minimum(xgrid), maximum(xgrid))
    PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,0.0,minimum(xgrid), maximum(xgrid)) # no zeros
    PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
end

savedir = "$(pwd())/results/" * sName *"/";
try mkdir(savedir) catch; end
#CSV.write(savedir * sName * "_IRF_YY_AggSh"*string(sh_id)*"_pmean.csv", DataFrame(YY_IRF,:auto))
#CSV.write(savedir * sName * "_IRF_PhatDens_AggSh"*string(sh_id)*"_pmean.csv", DataFrame(PhatDens_IRF,:auto))
CSV.write(savedir * sName * "_IRF_PhatDens_AggSh"*string(sh_id)*"_pmean_nozeros.csv", DataFrame(PhatDens_IRF,:auto))


#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------

println("")
println("Generating IRFs for posterior draws... ")
println("")

nsim = size(Bpdraw)[1]
n_subseq                 = floor(Int,nsim/n_every)
YY_IRF_uncertainty       = zeros(H+1,n,n_subseq)

# hh=1 is steady state
# hh=2 is period of impact

for pp = 1:n_subseq

    time_init_loop = time_ns();
    println(" Draw number:  $pp")
    println(" Remaining draws:  $(n_subseq-pp)")
    println(" Posterior draw number: $(pp*n_every)")

    L      = Sigtrpdraw[pp*n_every,:,:]
    Btilde = Bpdraw[pp*n_every,:,:]
    response = IRFredu(Btilde,L,H,sh_size,sh_id)
    YY_IRF_uncertainty[2:end,:,pp] = response'

    PhatDens_IRF   = zeros(H+1,length(xgrid));
    #Gini_IRF       = zeros(H+1,1)

    # compute density IRFs
    for hh = 1:H+1
        PhatDensCoef_hh    = coefRecover(YY_IRF_uncertainty[hh,n_agg+1:end,pp]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #    PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF_uncertainty[hh,3,pp]+mean_unrate,minimum(xgrid), maximum(xgrid))
        PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,0.0,minimum(xgrid), maximum(xgrid)) # no zeros
        PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
    end

#    CSV.write(savedir * sName * "_IRF_YY_AggSh"*string(sh_id)*"_" * string(pp) * ".csv", DataFrame(YY_IRF_uncertainty[:,:,pp]))
#    CSV.write(savedir * sName * "_IRF_PhatDens_AggSh" * string(sh_id) * "_" * string(pp) * ".csv", DataFrame(PhatDens_IRF))
    CSV.write(savedir * sName * "_IRF_PhatDens_AggSh" * string(sh_id) * "_" * string(pp) * "_nozeros.csv", DataFrame(PhatDens_IRF))

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println(" Elapsed time = $(time_loop)")
    println("")

end
